package trace

import (
	"context"
	"github.com/go-kratos/kratos/v2/errors"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/propagation"
	semconv "go.opentelemetry.io/otel/semconv/v1.21.0"
	"go.opentelemetry.io/otel/trace"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
	"jk-time/webook/pkg/grpcx/interceptors"
)

type InterceptorBuilder struct {
	p      propagation.TextMapPropagator
	tracer trace.Tracer
	interceptors.Builder
	serviceName string
}

func NewInterceptorBuilder(p propagation.TextMapPropagator, tracer trace.Tracer, serviceName string) *InterceptorBuilder {
	return &InterceptorBuilder{p: p, tracer: tracer, serviceName: serviceName}
}

func (b *InterceptorBuilder) BuildClient() grpc.UnaryClientInterceptor {
	propagator := b.p
	if propagator == nil {
		propagator = otel.GetTextMapPropagator()
	}

	tracer := b.tracer
	if tracer == nil {
		tracer = otel.Tracer("webook/pkg/grpcx/interceptors/trace/builder.go")
	}
	attrs := []attribute.KeyValue{
		semconv.RPCSystemKey.String("grpc"),
		attribute.Key("rpc.grpc.kind").String("unary"),
		attribute.Key("rpc.component").String("client"),
	}
	return func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn,
		invoker grpc.UnaryInvoker, opts ...grpc.CallOption) (err error) {
		newAttrs := append(attrs, semconv.RPCMethodKey.String(method),
			semconv.NetPeerNameKey.String(b.serviceName))
		ctx, span := tracer.Start(ctx, method, trace.WithAttributes(newAttrs...),
			trace.WithSpanKind(trace.SpanKindClient))
		ctx = Inject(ctx, propagator)
		defer func() {
			if err != nil {
				span.RecordError(err)
				if e := errors.FromError(err); e != nil {
					span.SetAttributes(semconv.RPCGRPCStatusCodeKey.Int64(int64(e.Code)))
				}
				span.SetStatus(codes.Error, err.Error())
			} else {
				span.SetStatus(codes.Ok, "OK")
			}
			span.End()
		}()
		return invoker(ctx, method, req, reply, cc, opts...)
	}
}
func (b *InterceptorBuilder) BuildServer() grpc.UnaryServerInterceptor {
	propagator := b.p
	if propagator == nil {
		propagator = otel.GetTextMapPropagator()
	}

	tracer := b.tracer
	if tracer == nil {
		tracer = otel.Tracer("webook/pkg/grpcx/interceptors/trace/builder.go")
	}

	attrs := []attribute.KeyValue{
		semconv.RPCSystemKey.String("grpc"),
		attribute.Key("rpc.grpc.kind").String("unary"),
		attribute.Key("rpc.component").String("client"),
	}
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo,
		handler grpc.UnaryHandler) (resp any, err error) {
		ctx = Extract(ctx, propagator)
		ctx, span := tracer.Start(ctx, info.FullMethod, trace.WithAttributes(attrs...),
			trace.WithSpanKind(trace.SpanKindServer))
		defer span.End()
		span.SetAttributes(
			semconv.RPCMethodKey.String(info.FullMethod),
			semconv.NetPeerNameKey.String(b.PeerName(ctx)),
			attribute.Key("net.peer.ip").String(b.PeerIP(ctx)),
		)
		defer func() {
			if err != nil {
				span.RecordError(err)
				if e := errors.FromError(err); e != nil {
					span.SetAttributes(semconv.RPCGRPCStatusCodeKey.Int64(int64(e.Code)))
				}
				span.SetStatus(codes.Error, err.Error())
			} else {
				span.SetStatus(codes.Ok, "OK")
			}
		}()
		return handler(ctx, req)
	}
}
func Inject(ctx context.Context, p propagation.TextMapPropagator) context.Context {
	// 先看ctx里面有没有元数据
	md, ok := metadata.FromOutgoingContext(ctx)
	if !ok {
		md = metadata.New(map[string]string{})
	}
	p.Inject(ctx, GrpcHeaderCarrier(md))
	return metadata.NewOutgoingContext(ctx, md)
}
func Extract(ctx context.Context, p propagation.TextMapPropagator) context.Context {
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		md = metadata.MD{}
	}

	//把md 注入到ctx中
	return p.Extract(ctx, GrpcHeaderCarrier(md))
}

// GrpcHeaderCarrier ...
type GrpcHeaderCarrier metadata.MD

func (md GrpcHeaderCarrier) Get(k string) string {
	vals := metadata.MD(md).Get(k)
	if len(vals) > 0 {
		return vals[0]
	}
	return ""
}
func (mc GrpcHeaderCarrier) Set(key string, value string) {
	metadata.MD(mc).Set(key, value)
}
func (mc GrpcHeaderCarrier) Keys() []string {
	keys := make([]string, 0, len(mc))
	for k := range metadata.MD(mc) {
		keys = append(keys, k)
	}
	return keys
}
